use crate::config::LlamaConfigJson;
use crate::tensor::Tensor;
use safetensors::{Dtype, SafeTensors};
use serde::Serialize;
use serde_json::{json, Result, Value};

pub struct LLamaParams<T: Serialize> {
    // token_id to embedding lookup table
    pub embedding_table: Tensor<T>, // (vocab_size, dim)
    // decoder layer
    pub rms_att_w: Vec<Tensor<T>>, // (hidden_size, ) x layers
    pub wq: Vec<Tensor<T>>,        // (n_heads * head_size, hidden_size) x layers
    pub wk: Vec<Tensor<T>>,        // (n_kv_heads * head_size, hidden_size) x layers
    pub wv: Vec<Tensor<T>>,        // (n_kv_heads * head_size, hidden_size) x layers
    pub wo: Vec<Tensor<T>>,        // (hidden_size, n_heads * head_size) x layers
    // ffn layer
    pub rms_ffn_w: Vec<Tensor<T>>, // (hidden_size, ) x layers
    pub w_up: Vec<Tensor<T>>,      // (intermediate_size, hidden_size) x layers
    pub w_gate: Vec<Tensor<T>>,    // (intermediate_size, hidden_size) x layers
    pub w_down: Vec<Tensor<T>>,    // (hidden_size, intermediate_size) x layers
    // output
    pub rms_out_w: Tensor<T>, // (hidden_size, )
    pub lm_head: Tensor<T>,   // (vocab_size, dim)
}

impl LLamaParams<f32> {
    pub fn from_safetensors(safetensor: &SafeTensors, config: &LlamaConfigJson) -> Self {
        let get_tensor = |name: &str| {
            let view = safetensor.tensor(name).unwrap();

            // ensure that the data-type is f32
            assert!(
                view.dtype() == Dtype::F32,
                "[ERROR] @from_safetensors >> data-type should be f32, but it is {:?}",
                view.dtype()
            );

            // ensure data is 4 bytes
            let capacity = view.data().len();
            assert!(
                capacity % 4 == 0,
                "[ERROR] @from_safetensors >> f32 should be 4 bytes, but the length of 1-byte data is {:?}",
                capacity
            );

            // convert u8 -> f32
            let data = view
                .data()
                .chunks(4)
                .map(|chunk| {
                    f32::from_le_bytes(chunk.try_into().expect("Chunk is not 4 bytes long"))
                })
                .collect();
            Tensor::<f32>::new(data, &Vec::from(view.shape()))
        };

        let layer = config.num_hidden_layers;

        LLamaParams {
            embedding_table: get_tensor("lm_head.weight"),

            rms_att_w: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.input_layernorm.weight")))
                .collect(),
            wq: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.self_attn.q_proj.weight")))
                .collect(),
            wk: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.self_attn.k_proj.weight")))
                .collect(),
            wv: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.self_attn.v_proj.weight")))
                .collect(),
            wo: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.self_attn.o_proj.weight")))
                .collect(),

            rms_ffn_w: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.post_attention_layernorm.weight")))
                .collect(),
            w_up: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.mlp.up_proj.weight")))
                .collect(),
            w_gate: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.mlp.gate_proj.weight")))
                .collect(),
            w_down: (0..layer)
                .map(|f| get_tensor(&format!("model.layers.{f}.mlp.down_proj.weight")))
                .collect(),

            rms_out_w: get_tensor("model.norm.weight"),
            lm_head: get_tensor("lm_head.weight"),
        }
    }

    pub fn to_json(&self) -> Result<Value> {
        let embedding_table_json = self.embedding_table.to_json()?;
        let rms_att_w_json: Result<Vec<Value>> =
            self.rms_att_w.iter().map(|t| t.to_json()).collect();
        let wq_json: Result<Vec<Value>> = self.wq.iter().map(|t| t.to_json()).collect();
        let wk_json: Result<Vec<Value>> = self.wk.iter().map(|t| t.to_json()).collect();
        let wv_json: Result<Vec<Value>> = self.wv.iter().map(|t| t.to_json()).collect();
        let wo_json: Result<Vec<Value>> = self.wo.iter().map(|t| t.to_json()).collect();
        let rms_ffn_w_json: Result<Vec<Value>> =
            self.rms_ffn_w.iter().map(|t| t.to_json()).collect();
        let w_up_json: Result<Vec<Value>> = self.w_up.iter().map(|t| t.to_json()).collect();
        let w_gate_json: Result<Vec<Value>> = self.w_gate.iter().map(|t| t.to_json()).collect();
        let w_down_json: Result<Vec<Value>> = self.w_down.iter().map(|t| t.to_json()).collect();
        let rms_out_w_json = self.rms_out_w.to_json()?;
        let lm_head_json = self.lm_head.to_json()?;

        Ok(json!({
            "embedding_table": embedding_table_json,
            "rms_att_w": rms_att_w_json?,
            "wq": wq_json?,
            "wk": wk_json?,
            "wv": wv_json?,
            "wo": wo_json?,
            "rms_ffn_w": rms_ffn_w_json?,
            "w_up": w_up_json?,
            "w_gate": w_gate_json?,
            "w_down": w_down_json?,
            "rms_out_w": rms_out_w_json,
            "lm_head": lm_head_json,
        }))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::use_example_dir;
    use std::fs::File;

    #[test]
    fn test_from_tensor() {
        use_example_dir(|model_dir| {
            let config = File::open(model_dir.join("config.json")).unwrap();
            let config: LlamaConfigJson = serde_json::from_reader(config).unwrap();
            let model_file = std::fs::read(model_dir.join("model.safetensors")).unwrap();
            let safetensor = SafeTensors::deserialize(&model_file).unwrap();
            let params = LLamaParams::from_safetensors(&safetensor, &config);

            assert!(params.to_json().is_ok());
        });
    }

    #[test]
    fn test_to_json() {
        use_example_dir(|model_dir| {
            let config = File::open(model_dir.join("config.json")).unwrap();
            let config: LlamaConfigJson = serde_json::from_reader(config).unwrap();
            let model_file = std::fs::read(model_dir.join("model.safetensors")).unwrap();
            let safetensor = SafeTensors::deserialize(&model_file).unwrap();
            let params = LLamaParams::from_safetensors(&safetensor, &config);

            assert!(params.to_json().is_ok());
        });
    }
}
